#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
局部键盘检测模块 - 检测按键相对位置关系
用于处理摄像头只能看到键盘局部的情况
"""

import cv2
import numpy as np
import numpy.typing as npt
from typing import Tuple, List, Dict, Optional, Any
import logging
from dataclasses import dataclass
from enum import Enum
import math

class KeyPosition(Enum):
    LEFT = "左侧"
    RIGHT = "右侧"
    UP = "上方"
    DOWN = "下方"
    DIAGONAL_UP_LEFT = "左上"
    DIAGONAL_UP_RIGHT = "右上"
    DIAGONAL_DOWN_LEFT = "左下"
    DIAGONAL_DOWN_RIGHT = "右下"

@dataclass
class KeyRelation:
    """按键关系"""
    key1: str
    key2: str
    position: KeyPosition
    distance: float
    angle: float
    confidence: float

@dataclass
class LocalDetectionResult:
    """局部检测结果"""
    detected_keys: List[Dict[str, Any]]
    key_relations: List[KeyRelation]
    sequence_errors: List[str]
    position_errors: List[str]
    confidence: float

class LocalKeyboardDetector:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
        # 标准键盘布局（QWERTY）
        self.standard_layout = {
            # 第一行
            '`': (0, 0), '1': (1, 0), '2': (2, 0), '3': (3, 0), '4': (4, 0), '5': (5, 0), '6': (6, 0), '7': (7, 0), '8': (8, 0), '9': (9, 0), '0': (10, 0), '-': (11, 0), '=': (12, 0),
            # 第二行
            'q': (0, 1), 'w': (1, 1), 'e': (2, 1), 'r': (3, 1), 't': (4, 1), 'y': (5, 1), 'u': (6, 1), 'i': (7, 1), 'o': (8, 1), 'p': (9, 1), '[': (10, 1), ']': (11, 1), '\\': (12, 1),
            # 第三行
            'a': (0, 2), 's': (1, 2), 'd': (2, 2), 'f': (3, 2), 'g': (4, 2), 'h': (5, 2), 'j': (6, 2), 'k': (7, 2), 'l': (8, 2), ';': (9, 2), "'": (10, 2),
            # 第四行
            'z': (0, 3), 'x': (1, 3), 'c': (2, 3), 'v': (3, 3), 'b': (4, 3), 'n': (5, 3), 'm': (6, 3), ',': (7, 3), '.': (8, 3), '/': (9, 3),
            # 空格键
            ' ': (4, 4)
        }
        
        # 按键大小估算
        self.estimated_key_size = 20  # 像素
        self.position_tolerance = 15  # 位置容差
        
    def detect_local_keyboard(self, image: npt.NDArray, reference_keys: List[Dict[str, Any]] = None) -> LocalDetectionResult:
        """
        检测局部键盘
        
        Args:
            image: 输入图像
            reference_keys: 参考按键列表（可选）
            
        Returns:
            LocalDetectionResult: 局部检测结果
        """
        # 检测图像中的按键
        detected_keys = self._detect_keys_in_image(image)
        
        if len(detected_keys) < 2:
            return LocalDetectionResult(
                detected_keys=detected_keys,
                key_relations=[],
                sequence_errors=[],
                position_errors=[],
                confidence=0.0
            )
        
        # 分析按键关系
        key_relations = self._analyze_key_relations(detected_keys)
        
        # 验证按键序列
        sequence_errors = self._validate_key_sequence(detected_keys)
        
        # 验证按键位置
        position_errors = self._validate_key_positions(detected_keys, key_relations)
        
        # 计算整体置信度
        confidence = self._calculate_confidence(detected_keys, key_relations, sequence_errors, position_errors)
        
        return LocalDetectionResult(
            detected_keys=detected_keys,
            key_relations=key_relations,
            sequence_errors=sequence_errors,
            position_errors=position_errors,
            confidence=confidence
        )
    
    def _detect_keys_in_image(self, image: npt.NDArray) -> List[Dict[str, Any]]:
        """检测图像中的按键"""
        # 转换为灰度图像
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
        
        # 预处理
        processed = self._preprocess_for_key_detection(gray)
        
        # 检测轮廓
        contours, _ = cv2.findContours(processed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        keys = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if 50 <= area <= 2000:  # 按键面积范围
                # 检查轮廓特征
                if self._is_valid_key_contour(contour):
                    key_info = self._extract_key_info(contour)
                    if key_info:
                        keys.append(key_info)
        
        return keys
    
    def _preprocess_for_key_detection(self, gray_image: npt.NDArray) -> npt.NDArray:
        """预处理图像用于按键检测"""
        # 自适应直方图均衡化
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray_image)
        
        # 高斯模糊
        blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
        
        # 自适应阈值
        binary = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, 11, 2
        )
        
        # 形态学操作
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        
        return cleaned
    
    def _is_valid_key_contour(self, contour: npt.NDArray) -> bool:
        """检查轮廓是否为有效的按键轮廓"""
        area = cv2.contourArea(contour)
        perimeter = cv2.arcLength(contour, True)
        
        if perimeter == 0:
            return False
        
        # 计算圆度
        circularity = 4 * np.pi * area / (perimeter * perimeter)
        
        # 计算长宽比
        x, y, w, h = cv2.boundingRect(contour)
        aspect_ratio = w / h if h > 0 else 0
        
        return (0.3 <= circularity <= 1.0 and 0.3 <= aspect_ratio <= 3.0)
    
    def _extract_key_info(self, contour: npt.NDArray) -> Optional[Dict[str, Any]]:
        """提取按键信息"""
        x, y, w, h = cv2.boundingRect(contour)
        center_x = x + w // 2
        center_y = y + h // 2
        
        return {
            'x': center_x,
            'y': center_y,
            'width': w,
            'height': h,
            'area': cv2.contourArea(contour),
            'contour': contour
        }
    
    def _analyze_key_relations(self, keys: List[Dict[str, Any]]) -> List[KeyRelation]:
        """分析按键之间的关系"""
        relations = []
        
        for i, key1 in enumerate(keys):
            for j, key2 in enumerate(keys):
                if i >= j:
                    continue
                
                # 计算距离和角度
                dx = key2['x'] - key1['x']
                dy = key2['y'] - key1['y']
                distance = math.sqrt(dx*dx + dy*dy)
                angle = math.degrees(math.atan2(dy, dx))
                
                # 确定相对位置
                position = self._determine_relative_position(dx, dy, angle)
                
                # 计算置信度
                confidence = self._calculate_relation_confidence(distance, angle, position)
                
                relation = KeyRelation(
                    key1=f"Key_{i}",
                    key2=f"Key_{j}",
                    position=position,
                    distance=distance,
                    angle=angle,
                    confidence=confidence
                )
                relations.append(relation)
        
        return relations
    
    def _determine_relative_position(self, dx: float, dy: float, angle: float) -> KeyPosition:
        """确定相对位置"""
        # 角度范围判断
        if -22.5 <= angle <= 22.5:
            return KeyPosition.RIGHT
        elif 22.5 < angle <= 67.5:
            return KeyPosition.DIAGONAL_DOWN_RIGHT
        elif 67.5 < angle <= 112.5:
            return KeyPosition.DOWN
        elif 112.5 < angle <= 157.5:
            return KeyPosition.DIAGONAL_DOWN_LEFT
        elif 157.5 < angle <= 180 or -180 <= angle <= -157.5:
            return KeyPosition.LEFT
        elif -157.5 < angle <= -112.5:
            return KeyPosition.DIAGONAL_UP_LEFT
        elif -112.5 < angle <= -67.5:
            return KeyPosition.UP
        elif -67.5 < angle <= -22.5:
            return KeyPosition.DIAGONAL_UP_RIGHT
        else:
            return KeyPosition.RIGHT
    
    def _calculate_relation_confidence(self, distance: float, angle: float, position: KeyPosition) -> float:
        """计算关系置信度"""
        # 基于距离的置信度
        distance_confidence = max(0, 1 - abs(distance - self.estimated_key_size) / self.estimated_key_size)
        
        # 基于角度的置信度（水平、垂直、对角线方向更可信）
        angle_confidence = 1.0
        if position in [KeyPosition.LEFT, KeyPosition.RIGHT]:
            angle_confidence = 1.0
        elif position in [KeyPosition.UP, KeyPosition.DOWN]:
            angle_confidence = 0.9
        else:
            angle_confidence = 0.8
        
        return (distance_confidence + angle_confidence) / 2
    
    def _validate_key_sequence(self, keys: List[Dict[str, Any]]) -> List[str]:
        """验证按键序列"""
        errors = []
        
        if len(keys) < 2:
            return errors
        
        # 按Y坐标排序（从上到下）
        sorted_keys = sorted(keys, key=lambda k: k['y'])
        
        # 检查每一行的按键顺序
        current_row = []
        current_y = sorted_keys[0]['y']
        
        for key in sorted_keys:
            if abs(key['y'] - current_y) <= self.position_tolerance:
                current_row.append(key)
            else:
                # 检查当前行的顺序
                if len(current_row) > 1:
                    row_errors = self._check_row_sequence(current_row)
                    errors.extend(row_errors)
                
                # 开始新行
                current_row = [key]
                current_y = key['y']
        
        # 检查最后一行的顺序
        if len(current_row) > 1:
            row_errors = self._check_row_sequence(current_row)
            errors.extend(row_errors)
        
        return errors
    
    def _check_row_sequence(self, row_keys: List[Dict[str, Any]]) -> List[str]:
        """检查单行按键顺序"""
        errors = []
        
        # 按X坐标排序（从左到右）
        sorted_row = sorted(row_keys, key=lambda k: k['x'])
        
        for i in range(len(sorted_row) - 1):
            key1 = sorted_row[i]
            key2 = sorted_row[i + 1]
            
            # 检查是否按正确顺序排列
            if key2['x'] <= key1['x']:
                errors.append(f"按键顺序错误: 位置 {key1['x']} 的按键应该在位置 {key2['x']} 的按键左侧")
        
        return errors
    
    def _validate_key_positions(self, keys: List[Dict[str, Any]], relations: List[KeyRelation]) -> List[str]:
        """验证按键位置"""
        errors = []
        
        # 检查按键间距是否合理
        for relation in relations:
            if relation.distance < self.estimated_key_size * 0.5:
                errors.append(f"按键距离过近: {relation.key1} 和 {relation.key2} 距离 {relation.distance:.1f} 像素")
            elif relation.distance > self.estimated_key_size * 3:
                errors.append(f"按键距离过远: {relation.key1} 和 {relation.key2} 距离 {relation.distance:.1f} 像素")
        
        return errors
    
    def _calculate_confidence(self, keys: List[Dict[str, Any]], relations: List[KeyRelation], 
                            sequence_errors: List[str], position_errors: List[str]) -> float:
        """计算整体置信度"""
        if not keys:
            return 0.0
        
        # 基础置信度
        base_confidence = 0.5
        
        # 按键数量奖励
        key_count_bonus = min(0.3, len(keys) * 0.05)
        
        # 关系质量奖励
        if relations:
            avg_relation_confidence = sum(r.confidence for r in relations) / len(relations)
            relation_bonus = avg_relation_confidence * 0.2
        else:
            relation_bonus = 0.0
        
        # 错误惩罚
        error_penalty = (len(sequence_errors) + len(position_errors)) * 0.05
        
        confidence = base_confidence + key_count_bonus + relation_bonus - error_penalty
        return max(0.0, min(1.0, confidence))
    
    def visualize_local_detection(self, image: npt.NDArray, result: LocalDetectionResult) -> npt.NDArray:
        """可视化局部检测结果"""
        vis_image = image.copy()
        
        # 绘制检测到的按键
        for i, key in enumerate(result.detected_keys):
            # 绘制按键边界框
            cv2.rectangle(
                vis_image,
                (key['x'] - key['width']//2, key['y'] - key['height']//2),
                (key['x'] + key['width']//2, key['y'] + key['height']//2),
                (0, 255, 0), 2
            )
            
            # 绘制按键编号
            cv2.putText(
                vis_image, f"K{i}", 
                (key['x'] - 10, key['y'] - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1
            )
        
        # 绘制按键关系线
        for relation in result.key_relations:
            key1_idx = int(relation.key1.split('_')[1])
            key2_idx = int(relation.key2.split('_')[1])
            
            if key1_idx < len(result.detected_keys) and key2_idx < len(result.detected_keys):
                key1 = result.detected_keys[key1_idx]
                key2 = result.detected_keys[key2_idx]
                
                # 绘制连接线
                color = (0, 255, 255) if relation.confidence > 0.7 else (0, 165, 255)
                cv2.line(vis_image, (key1['x'], key1['y']), (key2['x'], key2['y']), color, 1)
                
                # 绘制关系标签
                mid_x = (key1['x'] + key2['x']) // 2
                mid_y = (key1['y'] + key2['y']) // 2
                cv2.putText(
                    vis_image, relation.position.value[:2], 
                    (mid_x, mid_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1
                )
        
        # 绘制信息面板
        info_y = 30
        cv2.rectangle(vis_image, (10, 10), (350, 120), (0, 0, 0), -1)
        cv2.putText(vis_image, f"Detected Keys: {len(result.detected_keys)}", (20, info_y),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        info_y += 25
        cv2.putText(vis_image, f"Key Relations: {len(result.key_relations)}", (20, info_y),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        info_y += 25
        cv2.putText(vis_image, f"Sequence Errors: {len(result.sequence_errors)}", (20, info_y),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
        info_y += 25
        cv2.putText(vis_image, f"Confidence: {result.confidence:.2f}", (20, info_y),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        
        return vis_image
